// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use exitfailure::ExitFailure;
use failure::ResultExt;
use quicli::prelude::*;
use serde_derive::Deserialize;
use std::collections::BTreeMap;
use std::env;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use structopt::StructOpt;

const MESATEE_BUILD_CFG_DIR_ENV: &str = "MESATEE_BUILD_CFG_DIR";
const MESATEE_BUILD_CFG_FILE: &str = "build.toml";

#[derive(Debug, Deserialize)]
pub struct MesateeBuildToml {
    pub ra_config: BTreeMap<String, ConfigValue>,
    pub client_config: BTreeMap<String, ConfigValue>,
    pub audited_enclave_config: BTreeMap<String, ConfigValue>,
    pub rpc_config: BTreeMap<String, ConfigValue>,
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
pub enum ConfigValue {
    Simple(String),
    Numerical(u64),
    Full(ConfigComplex),
}

#[derive(Debug, Deserialize)]
pub struct ConfigComplex {
    pub path: Option<PathBuf>,
}

impl MesateeBuildToml {
    fn is_valid(&self) -> Result<(), ExitFailure> {
        let ra_config = &self.ra_config;
        if !ra_config.contains_key("mr_signer") {
            return Err(failure::err_msg("[ra_config]: missing `mr_signer`").into());
        }
        if !ra_config.contains_key("root_ca") {
            return Err(failure::err_msg("[ra_config]: missing `root_ca`").into());
        }
        if !ra_config.contains_key("ias_report_ca") {
            return Err(failure::err_msg("[ra_config]: missing `ias_report_ca`").into());
        }

        let client_config = &self.client_config;
        if !client_config.contains_key("client_cert") {
            return Err(failure::err_msg("[client_config]: missing `client_cert`").into());
        }
        if !client_config.contains_key("client_private_key") {
            return Err(failure::err_msg("[client_config]: missing `client_private_key`").into());
        }

        let audited_enclave_config = &self.audited_enclave_config;
        if !audited_enclave_config.contains_key("pubkey_a") {
            return Err(failure::err_msg("[audited_enclave_config]: missing `pubkey_a`").into());
        }
        if !audited_enclave_config.contains_key("pubkey_b") {
            return Err(failure::err_msg("[audited_enclave_config]: missing `pubkey_b`").into());
        }
        if !audited_enclave_config.contains_key("pubkey_c") {
            return Err(failure::err_msg("[audited_enclave_config]: missing `pubkey_c`").into());
        }

        let rpc_config = &self.rpc_config;
        if !rpc_config.contains_key("max_msg_size") {
            return Err(failure::err_msg("[rpc_config]: missing `max_msg_size`").into());
        }
        Ok(())
    }
}

fn generate_config_rs(
    cfg: &MesateeBuildToml,
    base_dir: &Path,
    dest_file: &Path,
) -> Result<(), ExitFailure> {
    let f = File::create(&dest_file)?;
    let mut writer = BufWriter::new(f);
    let _ = writer.write(b"//// This file is automatically generated. Do not modify.\n\n")?;
    let code = format!(
        r#"lazy_static! {{
    pub static ref MESATEE_SECURITY_CONSTANTS: MesateeSecurityConstants =
    MesateeSecurityConstants {{
        mr_signer: {},
        root_ca_bin: {},
        ias_report_ca: {},

        client_cert: {},
        client_pkcs8_key: {},

        audited_enclave_pubkey_a: {},
        audited_enclave_pubkey_b: {},
        audited_enclave_pubkey_c: {},

        max_msg_size: {:#x},
    }};

}}
"#,
        get_string(&cfg.ra_config["mr_signer"], &base_dir)?,
        get_string(&cfg.ra_config["root_ca"], &base_dir)?,
        get_string(&cfg.ra_config["ias_report_ca"], &base_dir)?,
        get_string(&cfg.client_config["client_cert"], &base_dir)?,
        get_string(&cfg.client_config["client_private_key"], &base_dir)?,
        get_string(&cfg.audited_enclave_config["pubkey_a"], &base_dir)?,
        get_string(&cfg.audited_enclave_config["pubkey_b"], &base_dir)?,
        get_string(&cfg.audited_enclave_config["pubkey_c"], &base_dir)?,
        get_numeric(&cfg.rpc_config["max_msg_size"])?
    );
    let _ = writer.write(code.as_bytes());
    writer.flush()?;
    Ok(())
}

fn get_string<P: AsRef<Path>>(
    config_value: &ConfigValue,
    base_dir: &P,
) -> Result<String, ExitFailure> {
    match config_value {
        ConfigValue::Simple(ref s) => Ok(s.to_string()),
        ConfigValue::Numerical(_) => {
            Err(failure::err_msg("Please specify a valid string".to_string()).into())
        }
        ConfigValue::Full(ref complex) => match &complex.path {
            Some(ref p) => {
                let path = base_dir.as_ref().to_path_buf().join(p);
                let include_fie_path = format!("include_bytes!({:?})", path);
                Ok(include_fie_path)
            }
            None => Err(failure::err_msg("Please specify a valid path".to_string()).into()),
        },
    }
}

fn get_numeric(config_value: &ConfigValue) -> Result<u64, ExitFailure> {
    match config_value {
        ConfigValue::Numerical(s) => Ok(*s),
        _ => Err(failure::err_msg("Please specify a valid number".to_string()).into()),
    }
}

fn load_mesatee_build_config_from_path(path: &Path) -> Result<MesateeBuildToml, ExitFailure> {
    let toml_data = read_file(path)?;
    let config: MesateeBuildToml = toml::from_str(&toml_data)?;
    config.is_valid()?;
    Ok(config)
}

#[derive(Debug, StructOpt)]
/// MesaTEE Configuration Generator
struct Cli {
    dest_file: String,
}

fn main() -> CliResult {
    let cfg_env =
        env::var(MESATEE_BUILD_CFG_DIR_ENV).with_context(|_| "$MESATEE_BUILD_CFG_DIR not set")?;
    let cfg_dir = Path::new(&cfg_env);
    if !cfg_dir.exists() || !cfg_dir.is_dir() {
        return Err(failure::err_msg(
            "$MESATEE_BUILD_CFG_DIR does not exist or is not a directory".to_string(),
        )
        .into());
    }
    let cfg_file = cfg_dir.to_path_buf().join(MESATEE_BUILD_CFG_FILE);
    if !cfg_file.exists() || !cfg_file.is_file() {
        return Err(failure::err_msg("build.toml does not exist".to_string()).into());
    }

    let config = load_mesatee_build_config_from_path(&cfg_file)?;
    let args = Cli::from_args();
    let dest_file = Path::new(&args.dest_file);
    generate_config_rs(&config, &cfg_dir, &dest_file)?;
    Ok(())
}
